import numpy as np

def predict_action(
    state: np.ndarray,
    lookback: int = 20,      # L：形成期（用过去 L 根计算收益）
    skip: int = 1,           # S：跳过最近 S 根，防止短期反转
    long_thresh: float = 0.01,   # 正向阈值（绝对收益或以波动率缩放后的值）
    short_thresh: float = -0.01, # 负向阈值
    use_vol_scale: bool = True,  # 是否用ATR做阈值缩放（自适应波动）
    atr_idx: int = 14,        # ATR 在 state 的列索引
    close_idx: int = 4,       # close 在 state 的列索引
):
    """
    单只择时（时间序列动量）策略（支持 batch：b >= 1）
    输入:
        state: np.ndarray, shape (b, 32, 20)
    参数:
        lookback, skip: 用于计算形成期动量
        long_thresh / short_thresh: 阈值，默认以绝对收益表示（2%）
        use_vol_scale: 若 True，则把阈值按 ATR 缩放：threshold *= atr / atr_ref
    输出:
        action: np.ndarray, shape (b,1), 值为 {-1,0,1}
    注意:
        - 要求 lookback + skip <= window-1 (window=32)
        - 当 ATR 无效或为0时退回到不缩放模式
    """
    b, window, c = state.shape
    assert window == 32 and c == 20, "期望 state 形状 (b,32,20)"
    assert lookback + skip <= window - 1, "lookback+skip 太长，超出窗口"

    # 索引
    idx_close = close_idx
    idx_atr = atr_idx

    # 时间点定位：最后一根为 t (index -1)
    # 我们计算 t-S 与 t-S-L 的价格
    idx_recent = -1 - skip            # 对应 P_{t-S}
    idx_past = -1 - skip - lookback   # 对应 P_{t-S-L}

    close_win = state[:, :, idx_close]   # shape (b, window)
    p_recent = close_win[:, idx_recent]   # (b,)
    p_past = close_win[:, idx_past]       # (b,)

    # 计算收益（可用简单收益或对数收益）
    # 小心除零和非有限值
    with np.errstate(divide='ignore', invalid='ignore'):
        mom = (p_recent / p_past) - 1.0   # (b,)
    mom = np.where(np.isfinite(mom), mom, np.nan)

    # 默认动作：观望
    action = np.zeros(b, dtype=np.int32)

    # 若全部无效，直接返回全 0
    valid_mask = np.isfinite(mom)
    if valid_mask.sum() == 0:
        return action.reshape(-1, 1)

    # 阈值自适应（基于 ATR）
    if use_vol_scale:
        atr_now = state[:, -1, idx_atr]      # (b,)
        # 避免 ATR<=0 导致异常：用小阈值替代
        atr_safe = np.where((atr_now > 0) & np.isfinite(atr_now), atr_now, np.nan)
        # 选一个参考 ATR（可用历史均值或样本中位数）。这里用样本中位数（有效值）
        ref_atr = np.nanmedian(atr_safe)
        if np.isfinite(ref_atr) and ref_atr > 0:
            scale = atr_safe / ref_atr    # 大于1表示更波动
            # 对有效元素才缩放阈值
            long_thresh_arr = np.full(b, long_thresh) * np.where(np.isfinite(scale), scale, 1.0)
            short_thresh_arr = np.full(b, short_thresh) * np.where(np.isfinite(scale), scale, 1.0)
        else:
            # 无法得到有效 ref_atr 时退回到常数阈值
            long_thresh_arr = np.full(b, long_thresh)
            short_thresh_arr = np.full(b, short_thresh)
    else:
        long_thresh_arr = np.full(b, long_thresh)
        short_thresh_arr = np.full(b, short_thresh)

    # 决策逻辑（对有效 mom 判定）
    # 多头： mom >= long_thresh_arr
    # 空头： mom <= short_thresh_arr
    # 否则观望
    long_idx = valid_mask & (mom >= long_thresh_arr)
    short_idx = valid_mask & (mom <= short_thresh_arr)

    action[long_idx] = 1
    action[short_idx] = -1

    return action.reshape(-1, 1)
